package org.apache.flink.kubernetes.kubeclient.decorators;

import org.apache.flink.kubernetes.kubeclient.ComponentPod;
import org.apache.flink.kubernetes.kubeclient.parameters.KubernetesProviderManagerParameters;
import org.apache.flink.kubernetes.utils.Constants;
import org.apache.flink.kubernetes.utils.KubernetesUtils;
import org.apache.flink.runtime.clusterframework.TaskExecutorProcessUtils;

import io.fabric8.kubernetes.api.model.Container;
import io.fabric8.kubernetes.api.model.ContainerBuilder;
import io.fabric8.kubernetes.api.model.EnvVarBuilder;

import java.util.List;

import static org.apache.flink.util.Preconditions.checkNotNull;

/** Attach the command and args to the main container for running the TaskManager. */
public class CmdProviderManagerDecorator extends AbstractKubernetesStepDecorator {

    private final KubernetesProviderManagerParameters kubernetesProviderManagerParameters;

    public CmdProviderManagerDecorator(
            KubernetesProviderManagerParameters kubernetesProviderManagerParameters) {
        this.kubernetesProviderManagerParameters = checkNotNull(kubernetesProviderManagerParameters);
    }

    @Override
    public ComponentPod decoratePod(ComponentPod flinkPod) {
        final Container mainContainerWithStartCmd =
                new ContainerBuilder(flinkPod.getMainContainer())
                        .withCommand(kubernetesProviderManagerParameters.getContainerEntrypoint())
                        .withArgs(getTaskManagerStartCommand())
                        .addToEnv(
                                new EnvVarBuilder()
                                        .withName(Constants.ENV_TM_JVM_MEM_OPTS)
                                        .withValue(
                                                kubernetesProviderManagerParameters.getJvmMemOptsEnv())
                                        .build())
                        .build();

        return new ComponentPod.Builder(flinkPod).withMainContainer(mainContainerWithStartCmd).build();
    }

    private List<String> getTaskManagerStartCommand() {
        final String resourceArgs =
                TaskExecutorProcessUtils.generateDynamicConfigsStr(
                        kubernetesProviderManagerParameters
                                .getContaineredTaskManagerParameters()
                                .getTaskExecutorProcessSpec());

        final String dynamicProperties = kubernetesProviderManagerParameters.getDynamicProperties();

        return KubernetesUtils.getStartCommandWithBashWrapper(
                Constants.KUBERNETES_PROVIDER_MANAGER_SCRIPT_PATH
                        + " "
                        + dynamicProperties
                        + " "
                        + resourceArgs
                        + " "
                        + kubernetesProviderManagerParameters.getEntrypointArgs());
    }
}
